{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6e3f0f72",
   "metadata": {},
   "source": [
    "# [beta] Structured Output\n",
    "\n",
    "It is often crucial to have LLMs return structured output. This is because often times the outputs of the LLMs are used in downstream applications, where specific arguments are required. Having the LLM return structured output reliably is necessary for that.\n",
    "\n",
    "There are a few different high level strategies that are used to do this:\n",
    "\n",
    "- Prompting: This is when you ask the LLM (very nicely) to return output in the desired format (JSON, XML). This is nice because works with all LLMs, this is not nice because it doesn't garuntee that the LLM returns in the right format.\n",
    "- Function calling: This is when the LLM is finetuned to be able to not just generate a completion, but also generate a function call. The functions the LLM can call are generally passed as extra parameters to the model API. The function names and descriptions should be treated as part of the prompt (they usually count against token counts, and are used by the LLM to decide what to do).\n",
    "- Tool calling: A technique similar to function calling, but it allows the LLM to call multiple functions at the same time.\n",
    "- JSON mode: This is when the LLM is garunteed to return JSON.\n",
    "\n",
    "\n",
    "\n",
    "Different models may support different variants of these, with slightly different parameters. In order to make it easy to get LLMs to return structured output, we have added a common interface to LangChain models: `.with_structured_output`. \n",
    "\n",
    "By invoking this method (and passing in a JSON schema or a Pydantic model) the model will add whatever model parameters + output parsers are necessary to get back the structured output. There may be more than one way to do this (eg function calling vs JSON mode) - you can configure which method to use by passing into that method.\n",
    "\n",
    "Let's look at some examples of this in action!\n",
    "\n",
    "We will use Pydantic to easily structure the response schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "08029f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.pydantic_v1 import BaseModel, Field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "070bf702",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Joke(BaseModel):\n",
    "    setup: str = Field(description=\"The setup of the joke\")\n",
    "    punchline: str = Field(description=\"The punchline to the joke\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98f6edfa",
   "metadata": {},
   "source": [
    "## OpenAI\n",
    "\n",
    "OpenAI exposes a few different ways to get structured outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3fe7caf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deddb6d3",
   "metadata": {},
   "source": [
    "### Function Calling\n",
    "\n",
    "By default, we will use `function_calling`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6700994a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatOpenAI()\n",
    "model_with_structure = model.with_structured_output(Joke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c55a61b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why was the cat sitting on the computer?', punchline='It wanted to keep an eye on the mouse!')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_with_structure.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39d7a555",
   "metadata": {},
   "source": [
    "### JSON Mode\n",
    "\n",
    "We also support JSON mode. Note that we need to specify in the prompt the format that it should respond in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "df0370e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_with_structure = model.with_structured_output(Joke, method=\"json_mode\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "23844a26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup=\"Why don't cats play poker in the jungle?\", punchline='Too many cheetahs!')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_with_structure.invoke(\n",
    "    \"Tell me a joke about cats, respond in JSON with `setup` and `punchline` keys\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f3cce9e",
   "metadata": {},
   "source": [
    "## Fireworks\n",
    "\n",
    "[Fireworks](https://fireworks.ai/) similarly supports function calling and JSON mode for select models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ad45fdd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_fireworks import ChatFireworks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36270ed5",
   "metadata": {},
   "source": [
    "### Function Calling\n",
    "\n",
    "By default, we will use `function_calling`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "49a20847",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatFireworks(model=\"accounts/fireworks/models/firefunction-v1\")\n",
    "model_with_structure = model.with_structured_output(Joke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e3093a6c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup=\"Why don't cats play poker in the jungle?\", punchline='Too many cheetahs!')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_with_structure.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb6b3ba",
   "metadata": {},
   "source": [
    "### JSON Mode\n",
    "\n",
    "We also support JSON mode. Note that we need to specify in the prompt the format that it should respond in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ea0c22c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_with_structure = model.with_structured_output(Joke, method=\"json_mode\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "649f9632",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why did the dog sit in the shade?', punchline='To avoid getting burned.')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_with_structure.invoke(\n",
    "    \"Tell me a joke about dogs, respond in JSON with `setup` and `punchline` keys\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff70609a",
   "metadata": {},
   "source": [
    "## Mistral\n",
    "\n",
    "We also support structured output with Mistral models, although we only support function calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bffd3fad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_mistralai import ChatMistralAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c8bd7549",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatMistralAI(model=\"mistral-large-latest\")\n",
    "model_with_structure = model.with_structured_output(Joke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17b15816",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_with_structure.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bbbb698",
   "metadata": {},
   "source": [
    "## Together\n",
    "\n",
    "Since [TogetherAI](https://www.together.ai/) is just a drop in replacement for OpenAI, we can just use the OpenAI integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9b9617e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "90549664",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatOpenAI(\n",
    "    base_url=\"https://api.together.xyz/v1\",\n",
    "    api_key=os.environ[\"TOGETHER_API_KEY\"],\n",
    "    model=\"mistralai/Mixtral-8x7B-Instruct-v0.1\",\n",
    ")\n",
    "model_with_structure = model.with_structured_output(Joke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "01da39be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why did the cat sit on the computer?', punchline='To keep an eye on the mouse!')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_with_structure.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3066b2af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
